import json
import os
import random
import numpy as np
import open3d as o3d
import tkinter as tk
from tkinter import ttk
from plyfile import PlyData, PlyElement
from pathlib import Path
from PointCloudRegistration import registration_point_clouds

class GroupInfo:
    def __init__(self, name, ply_geo, obj_geo, color):
        self.name = name
        self.ply_geo = ply_geo
        self.obj_geo = obj_geo
        self.color = color
        self.ply_visible = False
        self.obj_visible = False

def calculate_point_distances(source_pcd, target_pcd, threshold=0.01):
    """计算点云之间的距离，并标记超过阈值的点"""
    distances = source_pcd.compute_point_cloud_distance(target_pcd)
    # 将Open3D的DoubleVector转换为numpy数组
    distances = np.asarray(distances)
    colors = np.zeros((len(source_pcd.points), 3))
    # 设置默认颜色为灰色
    colors[:] = [0.5, 0.5, 0.5]
    # 将超过阈值的点标记为红色
    colors[distances > threshold] = [1, 0, 0]
    return colors

def read_ply(ply_path):
    ply_data = PlyData.read(ply_path)
    vertex_element = ply_data['vertex']
    property_names = [prop.name for prop in vertex_element.properties]
    vertex_data = []
    for v in vertex_element:
        row = []
        for prop in vertex_element.properties:
            val = v[prop.name]
            if 'float' in prop.val_dtype:
                formatted = f"{val:.6f}".rstrip('0').rstrip('.') if '.' in f"{val:.6f}" else f"{val}"
            else:
                formatted = str(val)
            row.append(formatted)
        vertex_data.append(row)
    return vertex_data, property_names

def save_registration_results(group_idx, transformation):
    cache_dir = Path("cache")
    cache_dir.mkdir(exist_ok=True)
    cache_file = cache_dir / "registration_results.json"

    if cache_file.exists():
        with open(cache_file, 'r') as f:
            cache = json.load(f)
    else:
        cache = {}

    cache[str(group_idx)] = transformation
    with open(cache_file, 'w') as f:
        json.dump(cache, f, indent=4)


def load_registration_results(group_idx):
    cache_file = Path("cache/registration_results.json")
    if cache_file.exists():
        with open(cache_file, 'r') as f:
            cache = json.load(f)
            return cache.get(str(group_idx))
    return None


def process_group(group, group_idx):
    print(f"处理组 {group_idx}")
    obj_path = "registered_ply/"+group["output_file"]
    ply_path = os.path.splitext(obj_path)[0] + ".ply"
    tx, ty, tz = group["transform"]["translate"]

    try:
        source_points_array, _ = read_ply(ply_path)
        source_pcd = o3d.geometry.PointCloud()
        source_pcd.points = o3d.utility.Vector3dVector(source_points_array)

        cached_T = load_registration_results(group_idx)
        if cached_T is not None:
            T = np.array(cached_T)
        else:
            is_small_cloud = True  # 可以修改为False测试大型点云处理模式

            # 为小型点云设置多次尝试和随机种子
            max_attempts = 15 if is_small_cloud else 5
            random_seed = 42 if is_small_cloud else None
            T = registration_point_clouds(ply_path, obj_path, is_small_cloud=is_small_cloud, max_attempts=max_attempts, random_seed=random_seed)
            save_registration_results(group_idx, T.tolist())

        source_pcd.transform(T)
        source_pcd.translate([tx, ty, tz])

        mesh = o3d.io.read_triangle_mesh(obj_path)
        if not mesh.has_vertices():
            raise ValueError("OBJ无顶点")
        mesh.translate([tx, ty, tz])
        mesh.compute_vertex_normals()
        mesh.paint_uniform_color([0.5, 0.5, 0.5])

        target_pcd = mesh.sample_points_uniformly(number_of_points=len(source_pcd.points))
        if target_pcd.has_points():
            colors = calculate_point_distances(source_pcd, target_pcd)
            source_pcd.colors = o3d.utility.Vector3dVector(colors)
        else:
            source_pcd.paint_uniform_color([0.5, 0.5, 0.5])

        return GroupInfo(f"Group {group_idx}", source_pcd, mesh, [0.5, 0.5, 0.5])

    except Exception as e:
        print(f"处理出错：{e}")
        return GroupInfo(f"Group {group_idx}", None, None, [0.5, 0.5, 0.5])

def merge_and_prepare_geometries(metadata_path):
    with open(metadata_path, 'r') as f:
        metadata = json.load(f)
    groups = []
    for idx, group in enumerate(metadata["groups"]):
        groups.append(process_group(group, idx))
    return [g for g in groups if g.ply_geo is not None or g.obj_geo is not None]

class VisualizationApp:
    def __init__(self, groups):
        self.groups = groups
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window(width=800, height=800)
        self.root = tk.Tk()
        self.root.title("Visualization Controller")
        
        # Create checkboxes
        self.vars = []
        for i, group in enumerate(self.groups):
            frame = ttk.Frame(self.root)
            frame.pack(fill='x', padx=5, pady=2)
            ply_var = tk.BooleanVar(value=group.ply_visible)
            obj_var = tk.BooleanVar(value=group.obj_visible)
            ttk.Checkbutton(frame, text="PLY", variable=ply_var,
                          command=lambda g=group, v=ply_var: self.toggle_visibility(g, 'ply', v)).pack(side='left')
            ttk.Checkbutton(frame, text="OBJ", variable=obj_var,
                          command=lambda g=group, v=obj_var: self.toggle_visibility(g, 'obj', v)).pack(side='left')
            ttk.Label(frame, text=group.name).pack(side='left')
            self.vars.append((ply_var, obj_var))

        # Setup Open3D initial geometries
        self.update_geometries()
        
        # Start update loop
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.root.after(10, self.update_visualization)
        self.root.mainloop()

    def toggle_visibility(self, group, geo_type, var):
        if geo_type == 'ply':
            group.ply_visible = var.get()
        else:
            group.obj_visible = var.get()
        self.update_geometries()

    def update_geometries(self):
        self.vis.clear_geometries()
        for group in self.groups:
            if group.ply_visible and group.ply_geo:
                self.vis.add_geometry(group.ply_geo)
            if group.obj_visible and group.obj_geo:
                self.vis.add_geometry(group.obj_geo)

    def update_visualization(self):
        if self.vis.poll_events():
            self.vis.update_renderer()
            self.root.after(10, self.update_visualization)
        else:
            self.on_close()

    def on_close(self):
        self.root.destroy()
        self.vis.destroy_window()
        os._exit(0)

if __name__ == "__main__":
    groups = merge_and_prepare_geometries("registered_ply/output_metadata.json")
    VisualizationApp(groups)